from mmengine.config import read_base
with read_base():
    from .mgam import *

from mgamdata.models.MedNeXt import MM_MedNext_Encoder, MM_MedNext_Decoder_3D
from mgamdata.models.FastSlow import RelativeSimilaritySelfSup, GapPredictor, VecAngConstraint, SimPairDiscriminator
from mgamdata.mm.mmseg_Dev3D import DiceLoss_3D


# val_evaluator=val_cfg=val_dataloader=multi_dataset_Ts=None
embed_dims = 128
block_counts = [2,2,2,2,2,2,2,2,2]
MedNeXt_checkpoint = False  # torch.checkpoint

model = dict(
    type = RelativeSimilaritySelfSup,
    momentum=FastSlow_Momentum,
    checkpoint_nir=nir_checkpoint,
    encoder = dict(
        type=MM_MedNext_Encoder,
        in_channels=in_channels, # type: ignore
        embed_dims=embed_dims,
        kernel_size=5,
        exp_r=2,
        block_counts = block_counts,
        dim="3d",
        use_checkpoint=MedNeXt_checkpoint,
        norm_type='group',
    ),
    neck=None,
    decoder = dict(
        type=MM_MedNext_Decoder_3D,
        embed_dims=embed_dims,
        kernel_size=5,
        exp_r=2,
        block_counts = block_counts,
        num_classes=embed_dims,
        out_channels=embed_dims,
        use_checkpoint=MedNeXt_checkpoint,
        deep_supervision=deep_supervision,
        norm_type='group',
        ignore_index=0, # 仅对train acc计算有效
        loss_gt_key='gt_sem_seg', # ["gt_sem_seg_one_hot", "gt_sem_seg"]
        loss_decode=dict(
            type=DiceLoss_3D, 
            use_sigmoid=False, 
            ignore_1st_index=False, 
            batch_z=4, 
            # NOTE Severe performance overhead when not being set to None.
            # NOTE Prefer using `ignore_1st_index`.
            # NOTE Invalid Class (Defaults to the last class) has been masked out during preprocess.
            ignore_index=None, 
        ), 
    ),
    gap_head=dict(
        type=GapPredictor, 
        in_channels=embed_dims, 
        dim='3d',
    ), 
    sim_head=dict(
        type=SimPairDiscriminator, 
        in_channels=embed_dims, 
        dim='3d', 
        view_size=size, 
        sub_view_size=sub_view_size, 
        loss_weight=1e-2,
    ),
    vec_head=dict(
        type=VecAngConstraint, 
        in_channels=embed_dims, 
        dim='3d', 
        num_views=num_views,
    ),
)